#include "LabelEquivalenceSeg.h"
#include <cuda_runtime.h>
#include <thrust\fill.h>
#include <thrust\device_ptr.h>
#include <iostream>

__global__ void boundaryDetectionKernel(bool* in, bool* boundary_device, int width, int height){
	int x = (blockIdx.x * blockDim.x + threadIdx.x);
	int y = (blockIdx.y * blockDim.y + threadIdx.y);

	boundary_device[x + y * width] = false;

	
	if(in[x + y * width]){
		boundary_device[x + y * width] = x == 0 || y == 0 || x == width - 1 || y == height - 1;


		boundary_device[x + y * width] = !boundary_device[x + y * width] && (
			(!in[(x - 1 >= 0 ? x - 1 : 0) + y * width]	&& in[(x + 1 < width ? x + 1 : width - 1) + y * width])		||
			(in[(x - 1 >= 0 ? x - 1 : 0) + y * width]	&& !in[(x + 1 < width ? x + 1 : width - 1) + y * width])	||
			(!in[x + (y - 1 >= 0 ? y - 1 : 0) * width]	&& in[x + (y + 1 < height ? y + 1 : height - 1) * width])	||
			(in[x + (y - 1 >= 0 ? y - 1 : 0) * width]	&& !in[x + (y + 1 < height ? y + 1 : height - 1) * width])
			);		
	} 	
}


void LabelEquivalenceSeg::labelImage(float4* in){
	thrust::fill(
		thrust::device_ptr<int>(counter),
		thrust::device_ptr<int>(counter) + width * height,
		0
	);

	n_label reset;
	reset.normal.x = 0.0;
	reset.normal.y = 0.0;
	reset.normal.z = 0.0;
	reset.normal.w = 0.0;
	reset.label = -1;

	thrust::fill(
		thrust::device_ptr<n_label>(img),
		thrust::device_ptr<n_label>(img) + width * height,
		reset;
	);

	initialize(in);

	for(int i = 0; i < 10; i++){
		scan();
		analysis();
	}

	//postProcess();

	cudaMemcpy(boundary_host, boundary_device, sizeof(bool) * width * height, cudaMemcpyDeviceToHost);
	cudaMemcpy(label, img, sizeof(int) * width * height, cudaMemcpyDeviceToHost);

}

__global__ void initKernel(float4* in, LabelEquivalenceSeg::n_label* img, int* ref, int width, int height){
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	int l = x + y * width;
	img[l].normal = in[l];
	if(abs(in[l].x) < 1.1){
		img[l].label = l;		
	} else {
		img[l].label = -1;

	}

	/*ref[l] = l;
	img[l] = in[l] ? l : -1;
	ref[l] = l;*/
}


void LabelEquivalenceSeg::initialize(float4* in){
	initKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(in, img, ref, width, height);
	
	/*boundaryDetectionKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(in, boundary_device, width, height);*/
}


__device__ bool compNormal(float4* a, float4* b){
	return 
		abs(cos(sqrt(a->normal.x * b->normal.x + a->normal.y * b->normal.y 
			+ a->normal.z * b->normal.z))) < 3.141592653 / 8.0
		&&
		abs(a->normal.w - b->normal.w) < 200.0;
}


__device__ int getMin(
LabelEquivalenceSeg::n_label* up, 
LabelEquivalenceSeg::n_label* left, 
LabelEquivalenceSeg::n_label* center, 
LabelEquivalenceSeg::n_label* right, 
LabelEquivalenceSeg::n_label* down){
	int c;
	c = up->label > -1 && compNormal(&up->normal, &center->normal) && up->label < center->label ? up->label : center->label;
	c = left->label > -1 && compNormal(&left->normal, &center->normal) && left->label < c ? left->label : c;
	c = right->label > -1 && compNormal(&right->normal, &center->normal) && right->label < c ? right->label : c;
	return down->label > -1 && compNormal(&down->normal, &center->normal) && down->label < c ? down->label : c;
}


__global__ void scanKernel(LabelEquivalenceSeg::n_label* img, int* ref, int width, int height){
	int x = (blockIdx.x * blockDim.x + threadIdx.x);
	int y = (blockIdx.y * blockDim.y + threadIdx.y);
	LabelEquivalenceSeg::n_label label1 = img[x + y * width];
	int label2;

	if(img[x + y * width] > -1){
		label2 = getMin(
			img[x + (y - 1 > 0 ? y - 1 : 0) * width], 
			img[(x - 1 > 0 ? x - 1 : 0) + y * width],
			img[x + y * width],
			img[(x + 1 < width ? x + 1 : width) + y * width],
			img[x + (y + 1 < height ? y + 1 : height) * width]);
		
		if(label2 < label1.label)
			atomicMin(&ref[label1.label], label2);
	}	
}

void LabelEquivalenceSeg::scan(){
	scanKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(img, ref, width, height);
}

__global__ void analysisKernel(int* ref, LabelEquivalenceSeg::n_label* img, int width, int height){
	int x = (blockIdx.x * blockDim.x + threadIdx.x);
	int y = (blockIdx.y * blockDim.y + threadIdx.y);

	int label = x + y * width;
		
	if(img[label].label == label){
		int current = ref[label];
		
		do{
			current = ref[current];
		} while(current != ref[current]);
		
		ref[label] = current;		
	}

	//Labeling phase
	if(img[label] > -1)
		img[label].label = ref[img[label]];
}



void LabelEquivalenceSeg::analysis(){
	analysisKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(ref, img, width, height);
}


__global__ void countKernel(int* img, int* counter, int width, int height){
	int x = (blockIdx.x * blockDim.x + threadIdx.x);
	int y = (blockIdx.y * blockDim.y + threadIdx.y);


	if(img[x + y * width] > -1)
		atomicAdd(&counter[img[x + y * width]], 1);
}

__global__ void dropSmallClustersKernel(int* img, int* counter, int width, int height){
	int x = (blockIdx.x * blockDim.x + threadIdx.x);
	int y = (blockIdx.y * blockDim.y + threadIdx.y);

	if(counter[img[x + y * width]] < 500)
		img[x + y * width] = -1;
}

void LabelEquivalenceSeg::postProcess(){
	countKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(img, counter, width, height);

	dropSmallClustersKernel<<<dim3(20, 20), dim3(32, 24)>>>
		(img, counter, width, height);
}